# 用于更新股票日线数据

import pandas as pd
from sqlalchemy import create_engine
import tushare as ts
from datetime import datetime, timedelta

# 创建数据库连接引擎
engine = create_engine('mysql+pymysql://root:@127.0.0.1:3306/stocks?charset=utf8mb4&use_unicode=1')

# 初始化 pro 接口
pro = ts.pro_api('d499c0996f2683e29f260b0e62df6b67e33c234b073d361736e18c87')

# 查询不同的 ts_code
query = "SELECT DISTINCT ts_code FROM stocks_daily"
ts_codes = pd.read_sql(query, con=engine)['ts_code'].tolist()

# 获取电脑当日时间并转换为 YYYYMMDD 格式
end_date = datetime.now().strftime('%Y%m%d')

# 遍历每个 ts_code
for ts_code in ts_codes:
    # 查询该 ts_code 的最新 trade_date
    query = f"SELECT MAX(trade_date) as latest_date FROM stocks_daily WHERE ts_code = '{ts_code}'"
    latest_date_df = pd.read_sql(query, con=engine)
    latest_date = latest_date_df['latest_date'].values[0]
    print(f"最新日期: {latest_date}")
    # 处理 latest_date 为 YYYYMMDD 格式的字符串
    start_date = (latest_date + timedelta(days=1)).strftime('%Y%m%d')

    # 检查start_date是否大于end_date（避免无效查询）
    if start_date > end_date:
        print(f"{ts_code} 已是最新数据，无需更新")
        continue

    # 拉取新数据
    try:
        df = pro.daily(
            ts_code=ts_code,
            start_date=start_date,
            end_date=end_date
        )
    except Exception as e:
        print(f"拉取 {ts_code} 数据时出现错误: {e}")
        continue

    # 确保列名和数据库表结构一致
    db_columns = ['ts_code', 'trade_date', 'open', 'high', 'low', 'close',
                  'pre_close', 'change', 'pct_chg', 'vol', 'amount']

    # 只保留需要的列
    df = df[db_columns] if set(db_columns).issubset(df.columns) else df

    # 重命名 change 字段为 chg_amo 以匹配数据库
    if 'change' in df.columns:
        df = df.rename(columns={'change': 'chg_amo'})

    # 将 trade_date 转换为数据库需要的格式
    if 'trade_date' in df.columns:
        df['trade_date'] = pd.to_datetime(df['trade_date'], format='%Y%m%d').dt.date

    # 将新数据写入 stocks_daily 表
    if not df.empty:
        try:
            # 先检查是否有重复数据（双重保险）
            existing_dates_query = f"SELECT trade_date FROM stocks_daily WHERE ts_code = '{ts_code}'"
            existing_dates = pd.read_sql(existing_dates_query, con=engine)['trade_date'].tolist()
            df = df[~df['trade_date'].isin(existing_dates)]

            if not df.empty:
                df.to_sql('stocks_daily', con=engine, if_exists='append', index=False)
                print(
                    f"成功为 {ts_code} 写入 {len(df)} 条新数据，日期范围: {df['trade_date'].min()} 到 {df['trade_date'].max()}")
            else:
                print(f"{ts_code} 没有需要更新的新数据")
        except Exception as e:
            print(f"为 {ts_code} 写入数据时出现错误: {e}")
    else:
        print(f"{ts_code} 在 {start_date} 到 {end_date} 期间没有新数据。")